package com.nulldev.util.security.encryptAPI_rw.impls.rsa;

import com.nulldev.util.data.Variables;
import com.nulldev.util.data.Sets.DataSet;
import com.nulldev.util.io.IOUtils;
import com.nulldev.util.security.encryptAPI_rw.api.ICipherType;
import com.nulldev.util.security.encryptAPI_rw.api.IPKPKBlockingCipher;
import com.nulldev.util.security.encryptAPI_rw.api.IPKPKKeyGenerator;
import com.nulldev.util.security.encryptAPI_rw.api.IPKPKStreamingCipher;

import static com.nulldev.util.manage.Macros.*;

import java.io.InputStream;
import java.io.OutputStream;
import java.security.KeyFactory;
import java.security.KeyPair;
import java.security.KeyPairGenerator;
import java.security.PrivateKey;
import java.security.PublicKey;
import java.security.spec.PKCS8EncodedKeySpec;
import java.security.spec.X509EncodedKeySpec;

import javax.crypto.Cipher;

public class RSACipher implements IPKPKBlockingCipher, IPKPKStreamingCipher, IPKPKKeyGenerator {

	/**
	 * TODO and FIXME: - Loading keys is currently broken, no clue why...
	 */

	private static final String algorithm = "RSA";
	private static final String[] CIPHER_NAMES =
		{ "RSA", "RSA-2048", "RSA-4096" };

	private final int keySize;

	public RSACipher(final int keySize) {
		__assert(keySize < 1024 || keySize > 8192, "Invalid keySize!");
		__assert(keySize % 16 != 0, "keySize MUST be a multiple of 16!");
		this.keySize = keySize;
	}

	public int getKeySize() {
		return this.keySize;
	}

	@Override
	public String getEncryptorID() {
		return "ecr_v3.jre.rsa";
	}

	@Override
	public String[] getCipherNames() {
		return CIPHER_NAMES;
	}

	@Override
	public byte[] encrypt(final byte[] data, final byte[] privateKey, final byte[] publicKey) {
		try {
			final Cipher cipher = Cipher.getInstance(algorithm);
			cipher.init(Cipher.ENCRYPT_MODE, __make_publicKey(publicKey));
			return cipher.doFinal(data);
		} catch (Exception ex) {
			throw new RuntimeException(ex);
		}
	}

	private PublicKey __make_publicKey(final byte[] publicKey) {
		try {
			return KeyFactory.getInstance(algorithm).generatePublic(new X509EncodedKeySpec(publicKey));
		} catch (Exception e) {
			throw new UnsupportedOperationException(e);
		}
	}

	private PrivateKey __make_privateKey(final byte[] privateKey) {
		try {
			return KeyFactory.getInstance(algorithm).generatePrivate(new PKCS8EncodedKeySpec(privateKey));
		} catch (Exception e) {
			throw new UnsupportedOperationException(e);
		}
	}

	@Override
	public byte[] decrypt(final byte[] data, final byte[] privateKey, final byte[] publicKey) {
		try {
			final Cipher cipher = Cipher.getInstance(algorithm);
			cipher.init(Cipher.DECRYPT_MODE, __make_privateKey(privateKey));
			return cipher.doFinal(data);
		} catch (Exception ex) {
			throw new RuntimeException(ex);
		}
	}

	@Override
	public PKPKSet makeKey(final int keySize) {
		__assert(keySize < 1024 || keySize > 8192, "Invalid keySize!");
		__assert(keySize % 16 != 0, "keySize MUST be a multiple of 16!");
		try {
			final KeyPairGenerator keyGen = KeyPairGenerator.getInstance(algorithm);
			keyGen.initialize(keySize);

			final KeyPair key = keyGen.generateKeyPair();

			return new PKPKSet(key.getPublic().getEncoded(), key.getPrivate().getEncoded());
		} catch (Exception e) {
			throw new RuntimeException(e);
		}
	}

	@Override
	public DataSet<Object, Object> makeNativeKey(final int keySize) {
		__assert(keySize < 1024 || keySize > 8192, "Invalid keySize!");
		__assert(keySize % 16 != 0, "keySize MUST be a multiple of 16!");
		try {
			final KeyPairGenerator keyGen = KeyPairGenerator.getInstance(algorithm);
			keyGen.initialize(keySize);

			final KeyPair key = keyGen.generateKeyPair();

			return new DataSet<Object, Object>(key.getPublic(), key.getPrivate());
		} catch (Exception e) {
			throw new RuntimeException(e);
		}
	}

	@Override
	public long encrypt(final InputStream is, final OutputStream os, final byte[] privateKey, final byte[] publicKey) {
		Variables.requireNonNull(is, "is");
		Variables.requireNonNull(os, "os");
		long written = 0;
		try {
			final Cipher cipher = Cipher.getInstance(algorithm);
			cipher.init(Cipher.ENCRYPT_MODE, __make_publicKey(privateKey));
			final byte[] buf = new byte[IOUtils.MEMORY_ALLOC_BUFFER];
			int read = 0;
			while ((read = is.read(buf)) != -1) {
				os.write(cipher.doFinal(buf, 0, read));
				written += read;
			}
		} catch (Exception ex) {
			throw new RuntimeException(ex);
		}
		return written;
	}

	@Override
	public long decrypt(final InputStream is, final OutputStream os, final byte[] privateKey, final byte[] publicKey) {
		Variables.requireNonNull(is, "is");
		Variables.requireNonNull(os, "os");
		long written = 0;
		try {
			final Cipher cipher = Cipher.getInstance(algorithm);
			cipher.init(Cipher.DECRYPT_MODE, __make_privateKey(privateKey));
			final byte[] buf = new byte[IOUtils.MEMORY_ALLOC_BUFFER];
			int read = 0;
			while ((read = is.read(buf)) != -1) {
				os.write(cipher.doFinal(buf, 0, read));
				written += read;
			}
		} catch (Exception ex) {
			throw new RuntimeException(ex);
		}
		return written;
	}

	@Override
	public ICipherType getCipherType() {
		return ICipherType.PRIVATE_PUBLIC_KEY;
	}

}
